'''
@Company: TWL
@Author: xue jian
@Email: xuejian@kanzhun.com
@Date: 2020-06-19 15:06:57
'''
# -*- coding: UTF-8 -*-

import sys
import numpy as np
buckets = ['C1_1', 'E1_1', 'F1_1', 'M1_1', 'M2_1', 'S1_2', 'SX_1', 'T1_2', 'T1_4', 'T2_1', 'X1_1', 'X2_1', 'Y1_2', 'C2_1', 'E2_1', 'F2_1', 'M1_2', 'S1_1', 'S2_1', 'T1_1', 'T1_3', 'T1_5', 'WM_1', 'X1_2', 'Y1_1', 'Y2_1', 'Q1', 'Q2', 'S0', 'GP_1']

def count_bucket_rate(w_f):
    count_bucket = {}
    count_sum = 0
    for bucket in buckets:
        count = 0
        for _, _ in enumerate(open(w_f + bucket + "/out", 'r')):
            count += 1
        count_bucket [bucket] = count
        count_sum += count
    for k, v in count_bucket.items():
        count_bucket[k] = v/float(count_sum)
    return count_bucket

def merge_dates_bucket(f_path, dates, save_path, bucket_num={}):
    if f_path[-1] != '/':
        f_path += '/'
    for date in dates:
        w_f = open(save_path + date + ".merge", 'w')
        date_path = f_path + date + "/"
        if bucket_num == {}:
            bucket_num = count_bucket_rate(date_path)
        file_dict = {}
        for bucket in buckets:
            p = date_path + bucket + "/out"
            file_dict[bucket] = open(p, 'r')
        tmp_sum_data = []
        for bucket in buckets:
            num = int(1000 * bucket_num[bucket])
            for i in range(num):
                tmp_sum_data.append(file_dict[bucket].readline())
        while len(tmp_sum_data) > 1:
            tmp_sum_data = np.asarray(tmp_sum_data)
            np.random.shuffle(tmp_sum_data)
            w_f.writelines(tmp_sum_data.tolist())
            tmp_sum_data = []
            for bucket in buckets:
                num = int(1000 * bucket_num[bucket])
                for i in range(num):
                    tmp_l = file_dict[bucket].readline()
                    if tmp_l == "":
                        break
                    tmp_sum_data.append(tmp_l)
        w_f.close()
        for bucket in buckets:
            file_dict[bucket].close()

if __name__ == '__main__':
    dates = [sys.argv[1]]
    print(dates)
    merge_dates_bucket("/data3/training_data/boss_rank/", dates, "/data2/training_data/boss_rank/")